import scipy
import torch
import soundfile as sf
from diffusers import StableAudioPipeline

repo_id = "stabilityai/stable-audio-open-1.0" # 线上
model_path = "E:\code\deepLearning\stable-audio\model\stable-audio-open-1.0" # 本地
pipe = StableAudioPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.float16) # 在这里修改本地调用/使用线上
pipe = pipe.to("cuda")

# 获取用户输入的prompt
prompt = input("请输入英文声音描述 (例如: 'The sound of a hammer hitting a wooden surface.'): ")
negative_prompt = "Low quality."

# 设置生成器的随机种子
generator = torch.Generator("cuda").manual_seed(0)

# 执行生成过程
audio = pipe(
    prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=200,
    # audio_end_in_s=45.0,
    num_waveforms_per_prompt=1,
    generator=generator,
).audios

# 保存第一个生成的音频
output = audio[0].T.float().cpu().numpy()
# 使用prompt的前20个字符作为文件名，避免特殊字符
safe_filename = ''.join(e for e in prompt[:40] if e.isalnum())
sf.write(f"{safe_filename}.wav", output, pipe.vae.sampling_rate)
print(f"音频已保存为: {safe_filename}.wav")